import pymysql
from dbutils.pooled_db import PooledDB
from typing import List, Dict, Union, Optional
from component.c_log import LOG

TAG = "c_mysql"
class MySQLDB:
    def __init__(
        self,
        host: str,
        user: str,
        password: str,
        database: str,
        port: int = 3306,
        charset: str = 'utf8mb4'
    ):
        """
        初始化数据库连接池
        :param host: 数据库地址
        :param user: 用户名
        :param password: 密码
        :param database: 数据库名称
        :param port: 端口号，默认3306
        :param charset: 字符编码，默认utf8mb4
        """
        self.pool = PooledDB(
            creator=pymysql,
            host=host,
            user=user,
            password=password,
            database=database,
            port=port,
            charset=charset,
            cursorclass=pymysql.cursors.DictCursor,
            mincached=5,
            maxconnections=20,
            blocking=True  # 连接池耗尽时阻塞等待
        )
        self.database = database  # 用于检查表存在性

    def _execute(
        self,
        sql: str,
        params: Optional[Union[tuple, dict]] = None,
        is_query: bool = True
    ) -> Optional[Union[List[Dict], int]]:
        """
        内部执行方法（统一处理连接获取和事务）
        :param sql: SQL语句
        :param params: 参数
        :param is_query: 是否为查询操作（决定是否fetch结果）
        :return: 查询结果或影响行数
        """
        connection = self.pool.connection()
        cursor = None
        try:
            cursor = connection.cursor()
            affected_rows = cursor.execute(sql, params)
            if is_query:
                result = cursor.fetchall()
                connection.commit()
                return result
            else:
                connection.commit()
                return affected_rows
        except Exception as e:
            connection.rollback()
            LOG(TAG,f"执行SQL失败: {e}")
            return None if is_query else 0
        finally:
            if cursor:
                cursor.close()
            connection.close()

    def execute_query(
        self,
        sql: str,
        params: Optional[Union[tuple, dict]] = None
    ) -> Optional[List[Dict]]:
        """
        执行SQL查询并返回结果
        """
        return self._execute(sql, params, is_query=True)

    def execute_non_query(
        self,
        sql: str,
        params: Optional[Union[tuple, dict]] = None
    ) -> int:
        """
        执行无结果集的SQL（如INSERT/UPDATE/DELETE）
        """
        return self._execute(sql, params, is_query=False) or 0

    def insert(
        self,
        table: str,
        data: Dict
    ) -> Optional[int]:
        """
        插入单条数据
        """
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        connection = self.pool.connection()
        try:
            with connection.cursor() as cursor:
                cursor.execute(sql, tuple(data.values()))
                lastrowid = cursor.lastrowid
            connection.commit()
            return lastrowid
        except Exception as e:
            connection.rollback()
            LOG(TAG,f"插入数据失败: {e}")
            return None
        finally:
            connection.close()

    def bulk_insert(
        self,
        table: str,
        data_list: List[Dict]
    ) -> int:
        """
        批量插入数据
        """
        if not data_list:
            return 0

        first_keys = set(data_list[0].keys())
        for idx, record in enumerate(data_list[1:], 1):
            if set(record.keys()) != first_keys:
                raise ValueError(f"第{idx+1}条记录的字段与第一条不一致")

        columns = ', '.join(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(data_list[0]))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        values = [tuple(record.values()) for record in data_list]

        connection = self.pool.connection()
        try:
            with connection.cursor() as cursor:
                affected_rows = cursor.executemany(sql, values)
            connection.commit()
            return affected_rows
        except Exception as e:
            connection.rollback()
            LOG(TAG,f"批量插入失败: {e}")
            return 0
        finally:
            connection.close()

    def delete(
        self,
        table: str,
        condition: str,
        params: Optional[Union[tuple, dict]] = None
    ) -> int:
        """
        删除数据
        """
        sql = f"DELETE FROM {table} WHERE {condition}"
        return self.execute_non_query(sql, params)

    def update(
        self,
        table: str,
        data: Dict,
        condition: str,
        params: Optional[Union[tuple, dict]] = None
    ) -> int:
        """
        更新数据
        """
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        sql = f"UPDATE {table} SET {set_clause} WHERE {condition}"
        full_params = tuple(data.values()) + (params if isinstance(params, tuple) else ())
        return self.execute_non_query(sql, full_params)

    def select(
        self,
        table: str,
        columns: List[str] = ['*'],
        condition: Optional[str] = None,
        params: Optional[Union[tuple, dict]] = None,
        extra: Optional[str] = None
    ) -> Optional[List[Dict]]:
        """
        查询数据
        """
        columns_str = ', '.join(columns)
        sql = f"SELECT {columns_str} FROM {table}"
        if condition:
            sql += f" WHERE {condition}"
        if extra:
            sql += f" {extra}"
        return self.execute_query(sql, params)

    def table_exists(self, table_name: str) -> bool:
        """
        检查表是否存在
        """
        sql = """
        SELECT 1 
        FROM information_schema.tables 
        WHERE table_schema = %s 
          AND table_name = %s
        LIMIT 1
        """
        result = self.execute_query(sql, (self.database, table_name))
        return bool(result)

    def close(self) -> None:
        """关闭连接池（实际会等待所有连接归还）"""
        self.pool.close()
